from mongoengine import *
from user import *
import numpy as np
import pickle

BUFFER_SIZE = 100
NO_USER_ERR = 'User does not exist'
DUP_USER_ERR = 'User already exists'
INV_TRAINING_ID_ERR = 'Invalid training id'
DB_NAME = 'webservice1'
DICT_KEY='value'

class DUser(Document):
	userid = StringField(required=True, primary_key=True)
	gender = StringField()
	interestedin = StringField()
	methodid = StringField()
	
class Image(Document):
	default = BooleanField()
	orig_image = FileField()
	nn_compressed = DictField()
	user = ReferenceField(DUser, reverse_delete_rule=CASCADE)
	
class NNMethod(Document):
	strategy = FileField()
	user = ReferenceField(DUser, reverse_delete_rule=CASCADE)
	
class NNPCAMethod(Document):
	strategy = FileField()
	user = ReferenceField(DUser, reverse_delete_rule=CASCADE)
	
class TrainingImage(Document):
	trainingid = StringField(required=True, primary_key=True)
	image = FileField()
	nn_compressed = DictField()
	
class Database:

	def __init__(self):
		self._buffer = dict()
		self._buffer_key_list = []
		self._buffer_size = BUFFER_SIZE
		self._inactive_users_keys = set([])
		connect(DB_NAME)
		self._init_inactive_users_keys()
		
	def _init_inactive_users_keys(self):
		for user in DUser.objects:
			self._inactive_users_keys.add(user.userid)
		
	def _write_documents_to_db(self, duser, image, nn, nn_pca, user):
		duser.gender = user.get_gender()
		duser.interestedin = user.get_interestedin()
		duser.methodid = user.get_method()
		duser.save()
		
		image.default = user.get_default_image()
		image.orig_image = user.get_image()
		
		if not image.default:
			image.nn_compressed = {DICT_KEY:user.get_nn_compressed_image().\
								   tolist()}
		else:
			image.nn_compressed = None
			
		image.user = duser
		image.save()
		
		nn.strategy = pickle.dumps(user.get_nn_strategy(), protocol = 2)
		nn.user = duser
		nn.save()
		
		nn_pca.strategy = pickle.dumps(user.get_nn_pca_strategy(), protocol = 2)
		nn_pca.user = duser
		nn_pca.save()
		
	def _write_to_db(self, userid, user):
		duser = DUser.objects.get(userid=userid)
		image = Image.objects.get(user=duser)
		nn = NNMethod.objects.get(user=duser)
		nn_pca = NNPCAMethod.objects.get(user=duser)

		self._write_documents_to_db(duser, image, nn, nn_pca, user)
		
	def _read_from_db(self, userid):
		duser = DUser.objects.get(userid=userid)
		image = Image.objects.get(user=duser)
		nn = pickle.loads(NNMethod.objects.get(user=duser).strategy.read())
		nn.recover()
		nn_pca = pickle.loads(NNPCAMethod.objects.get(user=duser).\
							  strategy.read())
		
		gender = duser.gender
		interestedin = duser.interestedin
		methodid = duser.methodid
		default_image = image.default
		orig_image = image.orig_image.read()
		
		if not default_image:
			nn_compressed = np.array(image.nn_compressed[DICT_KEY])
		else:
			nn_compressed = None
		
		user = User(userid, gender, interestedin, nn, nn_pca, methodid, 
					default_image, orig_image, nn_compressed)
		
		return user
		
	def _delete_from_db(self, userid):
		duser = DUser.objects.get(userid=userid)
		duser.delete()
		
	def _add_to_buffer(self, userid, user):
		if not len(self._buffer) < self._buffer_size:
			removed_userid = self._buffer_key_list.pop(0)
			removed_user = self._buffer.pop(removed_userid)
			self._inactive_users_keys.add(removed_userid)
			self._write_to_db(removed_userid, removed_user)
		self._buffer[userid] = user
		self._inactive_users_keys -= set([userid])
		self._buffer_key_list.append(userid)
		
	def _get_from_buffer(self, userid):
		if not userid in self._buffer:
			user = self._read_from_db(userid)
			self._add_to_buffer(userid, user)
		else:
			user = self._buffer[userid]
			
		return user
		
	def _remove_from_buffer(self, userid):
		if userid in self._buffer:
			self._buffer.pop(userid)
			self._buffer_key_list.remove(userid)
		else:
			self._inactive_users_keys.remove(userid)
		
	def save(self):
		for key, value in self._buffer.items():
			self._write_to_db(key, value)
			
	def training_images_added(self):
		return len(TrainingImage.objects) > 0
		
	def add_training_image(self, trainingid, training_image, nn_compressed):
		training_doc = TrainingImage(trainingid=trainingid,image=training_image)
		training_doc.nn_compressed = {DICT_KEY:nn_compressed.tolist()}
		training_doc.save()
		
	def get_training_image(self, trainingid):
		try:
			train_image = TrainingImage.objects.get(trainingid=trainingid)
			return train_image.image.read(), \
				   np.array(train_image.nn_compressed[DICT_KEY])
		except Exception:
			raise Exception(INV_TRAINING_ID_ERR)
		
	def add_user(self, user):
		userid = user.get_id()
		if len(DUser.objects(userid=userid)) != 0:
			raise Exception(DUP_USER_ERR)
				
		duser = DUser(userid=userid)
		image = Image()
		nn = NNMethod()
		nn_pca = NNPCAMethod()
		
		self._write_documents_to_db(duser, image, nn, nn_pca, user)
		self._add_to_buffer(userid, user)
		
	def delete_user(self, userid):
		try:
			self._delete_from_db(userid)
			self._remove_from_buffer(userid)
		except Exception:
			raise Exception(NO_USER_ERR)
		
	def get_user(self, userid, cached=True):
		try:
			if cached:
				return self._get_from_buffer(userid)
			else:
				return self._read_from_db(userid)
		except Exception:
			raise Exception(NO_USER_ERR)
			
	def get_active_users(self):
		return self._buffer
		
	def get_inactive_users_ids(self):
		return list(self._inactive_users_keys)


